from transformers import pipeline
from PIL import Image
# import requests
import torch

pipe = pipeline(
    "image-text-to-text",
    model="/Users/michael/codes_ai/medgemma-4b-it",
    torch_dtype=torch.bfloat16,
    device="mps",
)

# Image attribution: Stillwaterising, CC0, via Wikimedia Commons
# image_url = "https://upload.wikimedia.org/wikipedia/commons/c/c8/Chest_Xray_PA_3-8-2010.png"
# image = Image.open(requests.get(image_url, headers={"User-Agent": "example"}, stream=True).raw)

role_radio =  "You are an expert radiologist."
content_radio = "Describe this X-ray. Please use Chinese"

role_skin =  "你是一名皮肤科专家."
content_skin = "描述一下这张图片中的皮肤病变,并给出处方"

def analyze(image_url):
    image = Image.open(image_url)
    messages = [
        {
            "role": "system",
            "content": [{"type": "text", "text": role_skin}]
        },
        {
            "role": "user",
            "content": [
                {"type": "text", "text": content_skin},
                {"type": "image", "image": image},
            ]
        }
    ]
    output = pipe(text=messages, max_new_tokens=200)
    print(output[0]["generated_text"][-1]["content"])

# analyze("./pic/02.png")
# analyze("./pic/01.png")
# analyze("./pic/00.png")
analyze("./pic/med001-shi.jpg")
analyze("./pic/med002-han.jpg")
analyze("./pic/med003-xun.jpg")
analyze("./pic/med004-zhi.jpg")